import os
import operator
from dotenv import load_dotenv
from typing import Annotated, TypedDict

from langchain_core.messages import AnyMessage, HumanMessage
from langchain_core.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.checkpoint.memory import MemorySaver

# --- 1. 环境准备 ---
# 为了更好地演示多步调用，我们引入一个更强大的搜索工具 Tavily
# 并保留我们的计算器工具。
load_dotenv()

# a. 定义工具
# 注意：Tavily 工具需要设置 TAVILY_API_KEY 环境变量
# 如果您没有，可以继续使用 DuckDuckGo，但它可能不如 Tavily 稳定
try:
    search_tool = TavilySearchResults(max_results=2)
except ImportError:
    print("TavilySearchResults 未找到，请运行 'pip install tavily-python'")
    from langchain_community.tools import DuckDuckGoSearchRun
    search_tool = DuckDuckGoSearchRun()

@tool
def python_calculator(expression: str) -> str:
    """一个可以执行 Python 数学表达式的计算器。"""
    try:
        result = eval(expression, {"__builtins__": None}, {})
        return f"表达式 '{expression}' 的计算结果是: {result}"
    except Exception as e:
        return f"执行表达式 '{expression}' 时出错: {e}"

tools = [search_tool, python_calculator]

# --- 2. 状态与图的定义 (与之前相同) ---
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]

llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))
model_with_tools = llm.bind_tools(tools)

def agent_node(state: AgentState):
    print("\n---AGENT: 思考中...---")
    response = model_with_tools.invoke(state["messages"])
    return {"messages": [response]}

tool_node = ToolNode(tools)

workflow = StateGraph(AgentState)
workflow.add_node("agent", agent_node)
workflow.add_node("tools", tool_node)
workflow.add_conditional_edges("agent", tools_condition)
workflow.add_edge("tools", "agent")
workflow.set_entry_point("agent")

# --- 3. 编译图，设置中断点 ---
# 我们在 'tools' 节点前设置中断，这样每次工具调用都需要审批
checkpointer = MemorySaver()
app = workflow.compile(checkpointer=checkpointer, interrupt_before=["tools"])

# --- 4. 健壮的、可处理多次中断的循环审批逻辑 ---
if __name__ == "__main__":
    # 一个需要多步操作的问题：先搜索，再计算
    question = "LangChain 和 LangGraph 当前的 GitHub 星标数分别是多少？它们的总和是多少？"
    
    config = {"configurable": {"thread_id": "user_session_2"}}
    inputs = {"messages": [HumanMessage(content=question)]}
    
    # 启动图的执行
    app.invoke(inputs, config)

    # 进入循环，直到图执行完毕
    while True:
        # 获取当前的状态快照
        snapshot = app.get_state(config)
        
        # 检查图是否已暂停，等待进入下一个节点
        if snapshot.next:
            print("\n---流程暂停，等待您的审批---")
            print(f"下一步计划执行的节点: {snapshot.next[0]}")

            # 特别处理工具节点的审批
            if snapshot.next[0] == "tools":
                last_message = snapshot.values["messages"][-1]
                print("工具调用详情:", last_message.tool_calls)
            
            user_approval = input("您是否批准执行下一步？ (yes/no/exit): ")

            if user_approval.lower() == "yes":
                print("\n---用户批准，继续执行---")
                # 从断点处恢复执行，让图继续运行到下一个中断点或结束
                app.invoke(None, config=config)
                # 使用 continue 回到循环的开始，重新检查图的状态
                continue
            elif user_approval.lower() == 'exit':
                print("\n---用户选择退出---")
                break
            else:
                print("\n---用户拒绝，流程终止---")
                break
        else:
            # 如果 snapshot.next 为空，说明图已经执行到了 END
            print("\n---流程已完成！---")
            final_message = snapshot.values["messages"][-1]
            print("最终答案:")
            final_message.pretty_print()
            break 